import os
import json
import inspect
import re
import ast
from typing import Dict, List, Any, Optional, Callable
from datetime import datetime
import importlib.util


class FunctionManager:
    """
    函数管理器 - 管理所有指标函数
    负责函数注册、调用、元信息提取等功能
    """
    
    def __init__(self, registry_file_path: str, functions_dir: str = None):
        self.registry_file = registry_file_path
        self.functions_dir = functions_dir or os.path.join(os.path.dirname(registry_file_path), "functions")
        self.functions = self._load_registry()
        self._function_cache = {}  # 缓存已加载的函数
        
        # 确保functions目录存在
        os.makedirs(self.functions_dir, exist_ok=True)
        # 确保__init__.py文件存在
        init_file = os.path.join(self.functions_dir, "__init__.py")
        if not os.path.exists(init_file):
            with open(init_file, 'w', encoding='utf-8') as f:
                f.write("# Functions directory\n")
    
    def _load_registry(self) -> Dict:
        """从文件加载函数注册表"""
        if os.path.exists(self.registry_file):
            try:
                with open(self.registry_file, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    # 只返回functions部分，不包含metadata
                    return data.get('functions', {})
            except Exception as e:
                print(f"Error loading function registry: {e}")
                return {}
        return {}
    
    def _save_registry(self):
        """保存函数注册表到文件"""
        try:
            os.makedirs(os.path.dirname(self.registry_file), exist_ok=True)
            
            # 构建完整的注册表结构
            registry_data = {
                "functions": self.functions,
                "metadata": {
                    "count": len(self.functions),
                    "last_updated": datetime.now().isoformat() if self.functions else None
                }
            }
            
            with open(self.registry_file, 'w', encoding='utf-8') as f:
                json.dump(registry_data, f, ensure_ascii=False, indent=2)
        except Exception as e:
            print(f"Error saving function registry: {e}")
    
    def _sanitize_filename(self, name: str) -> str:
        """将函数名称转换为有效的文件名"""
        import re
        sanitized = re.sub(r'\W|^(?=\d)', '_', name).lower()
        return sanitized.strip()
    
    def _get_unique_filename(self, indicator_name: str) -> str:
        """获取唯一的文件名"""
        base_filename = self._sanitize_filename(indicator_name)
        filename = f"{base_filename}.py"
        
        # 对于已注册的指标，始终使用相同的文件名
        if indicator_name in self.functions:
            return filename
            
        # 对于新指标，检查文件名是否被占用
        file_path = os.path.join(self.functions_dir, filename)
        if not os.path.exists(file_path):
            return filename
            
        # 有冲突，添加后缀
        counter = 1
        while True:
            new_filename = f"{base_filename}_{counter}.py"
            new_file_path = os.path.join(self.functions_dir, new_filename)
            if not os.path.exists(new_file_path):
                return new_filename
            counter += 1
    
    def _create_function_file(self, indicator_name: str, code: str):
        """创建单独的函数文件"""
        try:
            function_info = self.functions.get(indicator_name, {})
            filename = function_info.get('filename') or self._get_unique_filename(indicator_name)
            file_path = os.path.join(self.functions_dir, filename)
            
            with open(file_path, 'w', encoding='utf-8') as f:
                f.write(code)
            
            print(f"Created/Updated function file: {file_path}")
            
        except Exception as e:
            print(f"Error creating function file for {indicator_name}: {e}")
    
    def _extract_function_signature(self, code: str) -> Dict[str, Any]:
        """从代码中提取函数签名信息"""
        try:
            tree = ast.parse(code)
            
            for node in ast.walk(tree):
                if isinstance(node, ast.FunctionDef):
                    function_info = {
                        "name": node.name,
                        "parameters": [],
                        "return_annotation": None,
                        "docstring": ast.get_docstring(node)
                    }
                    
                    # 提取参数信息
                    for arg in node.args.args:
                        param_info = {
                            "name": arg.arg,
                            "annotation": None,
                            "default": None,
                            "required": True,
                            "description": ""
                        }
                        
                        # 提取类型注解
                        if arg.annotation:
                            param_info["annotation"] = ast.unparse(arg.annotation)
                        
                        function_info["parameters"].append(param_info)
                    
                    # 提取默认值
                    defaults = node.args.defaults
                    if defaults:
                        # 默认值从右往左对应参数
                        start_idx = len(function_info["parameters"]) - len(defaults)
                        for i, default in enumerate(defaults):
                            param_idx = start_idx + i
                            if param_idx < len(function_info["parameters"]):
                                function_info["parameters"][param_idx]["default"] = ast.unparse(default)
                                function_info["parameters"][param_idx]["required"] = False
                    
                    # 提取返回类型注解
                    if node.returns:
                        function_info["return_annotation"] = ast.unparse(node.returns)
                    
                    return function_info
                    
        except Exception as e:
            print(f"Error extracting function signature: {e}")
            return {}
        
        return {}
    
    def register_function(self, indicator_name: str, code: str, source_file: str = None, 
                         excel_created_at: str = None, excel_updated_at: str = None) -> str:
        """
        注册一个函数
        
        Args:
            indicator_name: 指标名称
            code: 生成的函数代码
            source_file: 源Excel文件名（可选）
            excel_created_at: Excel中的创建日期
            excel_updated_at: Excel中的更新日期
        
        Returns:
            str: 操作类型 ('new', 'updated', 'skipped')
        """
        existing_function = self.functions.get(indicator_name)
        
        # 提取函数签名信息
        signature_info = self._extract_function_signature(code)
        
        # 如果函数不存在，直接新增
        if not existing_function:
            filename = self._get_unique_filename(indicator_name)
            self.functions[indicator_name] = {
                'code': code,
                'source_file': source_file,
                'created_at': excel_created_at or datetime.now().isoformat(),
                'updated_at': excel_updated_at or datetime.now().isoformat(),
                'excel_created_at': excel_created_at,
                'excel_updated_at': excel_updated_at,
                'filename': filename,
                'signature': signature_info
            }
            self._save_registry()
            self._create_function_file(indicator_name, code)
            # 清除缓存
            if indicator_name in self._function_cache:
                del self._function_cache[indicator_name]
            return 'new'
        
        # 检查是否需要更新
        should_update = (
            existing_function.get('code', '') != code or
            source_file == 'manual_edit' or
            (excel_updated_at and existing_function.get('excel_updated_at') != excel_updated_at)
        )
        
        if should_update:
            filename = existing_function.get('filename') or self._get_unique_filename(indicator_name)
            self.functions[indicator_name] = {
                'code': code,
                'source_file': source_file,
                'created_at': existing_function.get('created_at', datetime.now().isoformat()),
                'updated_at': datetime.now().isoformat(),
                'excel_created_at': excel_created_at or existing_function.get('excel_created_at'),
                'excel_updated_at': excel_updated_at or existing_function.get('excel_updated_at'),
                'filename': filename,
                'signature': signature_info
            }
            self._save_registry()
            self._create_function_file(indicator_name, code)
            # 清除缓存
            if indicator_name in self._function_cache:
                del self._function_cache[indicator_name]
            return 'updated'
        
        return 'skipped'
    
    def get_function_list(self) -> List[Dict[str, Any]]:
        """
        获取所有函数的元信息列表，用于RAG知识库
        
        Returns:
            List[Dict]: 包含函数元信息的列表
        """
        function_list = []
        
        for indicator_name, function_info in self.functions.items():
            signature = function_info.get('signature', {})
            
            meta_info = {
                "function_name": indicator_name,
                "name": signature.get('name', indicator_name),
                "description": signature.get('docstring', f'{indicator_name}数据获取函数'),
                "parameters": [],
                "return_type": signature.get('return_annotation', 'dict'),
                "created_at": function_info.get('created_at'),
                "updated_at": function_info.get('updated_at'),
                "source_file": function_info.get('source_file')
            }
            
            # 处理参数信息
            for param in signature.get('parameters', []):
                param_name = param.get('name', '')
                param_type = param.get('annotation', 'Any')
                
                # 生成更有意义的参数描述
                description = param.get('description', '')
                if not description or description == f'{param_name}参数':
                    # 根据参数名和类型生成描述
                    if param_name == 'project_id':
                        description = '项目ID'
                    elif param_name == 'questionnaire_ids':
                        description = '问卷ID集合'
                    elif param_name == 'product_code':
                        description = '产品编码，用于路由分发'
                    elif param_name == 'project_code':
                        description = '项目编码，用于路由分发'
                    elif param_name == 'region_code':
                        description = '区域编码，用于路由分发'
                    elif param_name == 'education':
                        description = '学历筛选条件（如：本科毕业生、专科毕业生、硕士研究生、博士研究生）'
                    elif param_name == 'db':
                        description = '数据库连接对象（自动注入）'
                    else:
                        description = f'{param_name}参数'
                
                param_info = {
                    "name": param_name,
                    "type": param_type,
                    "required": param.get('required', True),
                    "default": param.get('default'),
                    "description": description
                }
                meta_info["parameters"].append(param_info)
            
            function_list.append(meta_info)
        
        return function_list
    
    def get_function_list_text(self) -> str:
        """
        获取函数列表的文本格式，用于RAG
        
        Returns:
            str: 格式化的函数列表文本，包含MCP调用信息
        """
        function_list = self.get_function_list()
        
        if not function_list:
            return "暂无可用函数"
        
        text_parts = [
            "=== MCP函数调用指南 ===\n",
            "所有函数都通过统一的MCP工具调用：call_indicator_function",
            "调用格式：call_indicator_function(function_name='函数名称', parameters={参数字典})",
            "注意：db参数会自动注入，无需在parameters中提供\n",
            "=== 可用函数列表 ===\n"
        ]
        
        for i, func in enumerate(function_list, 1):
            function_name = func['function_name']
            
            text_parts.append(f"{i}. 【{function_name}】")
            text_parts.append(f"   MCP调用: call_indicator_function(function_name='{function_name}', parameters={{...}})")
            text_parts.append(f"   功能描述: {func['description']}")
            text_parts.append(f"   返回类型: {func['return_type']}")
            
            # 构建参数信息
            required_params = []
            optional_params = []
            
            if func['parameters']:
                for param in func['parameters']:
                    if param['name'] == 'db':  # 跳过db参数，因为会自动注入
                        continue
                    
                    param_info = f"{param['name']} ({param['type']})"
                    if param.get('description'):
                        param_info += f" - {param['description']}"
                    
                    if param['required']:
                        required_params.append(param_info)
                    else:
                        default_val = param.get('default', 'None')
                        param_info += f" [默认: {default_val}]"
                        optional_params.append(param_info)
            
            if required_params:
                text_parts.append("   必需参数:")
                for param in required_params:
                    text_parts.append(f"     - {param}")
            
            if optional_params:
                text_parts.append("   可选参数:")
                for param in optional_params:
                    text_parts.append(f"     - {param}")
            
            # 添加调用示例
            example_params = {}
            if func['parameters']:
                for param in func['parameters']:
                    if param['name'] == 'db':
                        continue
                    if param['required']:
                        if 'int' in param['type']:
                            example_params[param['name']] = 12345
                        elif 'List[int]' in param['type']:
                            example_params[param['name']] = [1, 2, 3]
                        elif 'str' in param['type']:
                            example_params[param['name']] = f"示例{param['name']}"
                        else:
                            example_params[param['name']] = f"<{param['name']}>"
            
            if example_params:
                import json
                example_json = json.dumps(example_params, ensure_ascii=False, indent=2)
                text_parts.append(f"   调用示例:")
                text_parts.append(f"     call_indicator_function(")
                text_parts.append(f"         function_name='{function_name}',")
                text_parts.append(f"         parameters={example_json}")
                text_parts.append(f"     )")
            
            text_parts.append("")  # 空行分隔
        
        return "\n".join(text_parts)
    
    def load_function(self, indicator_name: str) -> Optional[Callable]:
        """
        动态加载并返回指定的函数
        
        Args:
            indicator_name: 指标名称
            
        Returns:
            Callable: 加载的函数对象，如果失败返回None
        """
        if indicator_name in self._function_cache:
            return self._function_cache[indicator_name]
        
        function_info = self.functions.get(indicator_name)
        if not function_info:
            print(f"Function {indicator_name} not found in registry")
            return None
        
        filename = function_info.get('filename')
        if not filename:
            print(f"No filename for function {indicator_name}")
            return None
        
        file_path = os.path.join(self.functions_dir, filename)
        if not os.path.exists(file_path):
            print(f"Function file not found: {file_path}")
            return None
        
        try:
            # 动态加载模块
            spec = importlib.util.spec_from_file_location(f"function_{indicator_name}", file_path)
            module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(module)
            
            # 查找函数
            signature = function_info.get('signature', {})
            func_name = signature.get('name')
            
            if func_name and hasattr(module, func_name):
                func = getattr(module, func_name)
                self._function_cache[indicator_name] = func
                return func
            else:
                print(f"Function {func_name} not found in module {file_path}")
                return None
                
        except Exception as e:
            print(f"Error loading function {indicator_name}: {e}")
            return None
    
    def call_function(self, function_name: str, **kwargs) -> Any:
        """
        调用指定的函数
        
        Args:
            function_name: 函数名称（指标名称）
            **kwargs: 函数参数
            
        Returns:
            Any: 函数执行结果
        """
        func = self.load_function(function_name)
        if not func:
            raise ValueError(f"Function {function_name} not available")
        
        try:
            # 获取函数签名，只传递函数实际需要的参数
            sig = inspect.signature(func)
            filtered_kwargs = {}
            
            for param_name, param in sig.parameters.items():
                if param_name in kwargs:
                    filtered_kwargs[param_name] = kwargs[param_name]
                elif param.default is not param.empty:
                    # 参数有默认值，不需要传递
                    continue
                else:
                    # 必需参数但未提供，保持原有行为
                    pass
            
            return func(**filtered_kwargs)
        except Exception as e:
            print(f"Error calling function {function_name}: {e}")
            raise
    
    def get_function_info(self, function_name: str = None) -> Dict:
        """
        获取函数信息
        
        Args:
            function_name: 函数名称，如果为None则返回所有函数信息
            
        Returns:
            Dict: 函数信息
        """
        if function_name:
            return self.functions.get(function_name, {})
        else:
            return {
                'total_functions': len(self.functions),
                'functions': self.functions
            }
    
    def remove_function(self, function_name: str) -> bool:
        """
        移除函数
        
        Args:
            function_name: 函数名称
            
        Returns:
            bool: 是否成功移除
        """
        if function_name not in self.functions:
            print(f"Function {function_name} not found in registry")
            return False
        
        try:
            # 删除文件
            function_info = self.functions[function_name]
            filename = function_info.get('filename')
            if filename:
                file_path = os.path.join(self.functions_dir, filename)
                if os.path.exists(file_path):
                    os.remove(file_path)
                    print(f"Removed function file: {file_path}")
            
            # 从注册表中移除
            del self.functions[function_name]
            self._save_registry()
            print(f"Removed function from registry: {function_name}")
            
            # 清除缓存
            if function_name in self._function_cache:
                del self._function_cache[function_name]
            
            return True
            
        except Exception as e:
            print(f"Error removing function {function_name}: {e}")
            return False
    
    def clear_all(self):
        """清空所有函数"""
        try:
            # 删除所有函数文件
            for function_name, function_info in self.functions.items():
                filename = function_info.get('filename')
                if filename:
                    file_path = os.path.join(self.functions_dir, filename)
                    if os.path.exists(file_path):
                        os.remove(file_path)
            
            # 清空注册表
            self.functions = {}
            self._save_registry()
            
            # 清空缓存
            self._function_cache = {}
            
            print("All functions cleared")
            
        except Exception as e:
            print(f"Error clearing functions: {e}")
    
    def get_function_count(self) -> int:
        """获取函数总数"""
        return len(self.functions)
